"""
多尺度互信息损失函数实现
包含三个主要损失分支：
1. CE Loss - 交叉熵损失
2. Global Information Loss - 全局信息损失
3. Local Information Loss - 局部互信息损失

- 支持自动权重平衡
- 可选KL散度或Wasserstein距离计算互信息
- 温度系数调节软化概率分布
"""

import logging
import torch
from losses import distance
import torch.nn as nn
from losses.AutomaticWeightedLoss import AutomaticWeightedLoss
import numpy as np

class loss_functions():
    def __init__(self, method='distance', mi_calculator='kl', temperature=1.5, bml_method='auto', scales=[1, 1, 1],
                 gil_loss=False,lil_loss=False,device='cuda:0'):
        """
        初始化损失函数类
        Args:
            method: 互信息计算方法，'distance'或'mi'
            mi_calculator: 互信息计算器类型，'kl'(KL散度)或'w'(Wasserstein距离)
            temperature: softmax温度系数，用于软化概率分布
            bml_method: 损失平衡方法，'auto'(自动权重)或'hyper'(超参数)
            scales: 各损失的权重系数
            gil_loss: 是否使用全局信息损失
            lil_loss: 是否使用局部信息损失
            device: 计算设备
        """
        self.lil_loss=lil_loss
        self.gil_loss=gil_loss

        # 确定损失函数数量
        loss_num=1  # 基础CE loss
        if gil_loss:
            loss_num+=1
            logging.info("With Global Information Loss")
        if lil_loss:
            loss_num+=1
            logging.info("With Local Information Loss")
        self.balance_loss = AutomaticWeightedLoss(loss_num)  # 初始化自动权重平衡器

        # 配置互信息计算相关参数
        self.softmax = torch.nn.Softmax(dim=1)
        self.method=method
        self.bml_method =bml_method
        self.scales=scales
        print(f"Mutual Information Calculator is :{mi_calculator}")
        
        # 选择互信息计算方法
        if mi_calculator == "kl":
            self.mi_calculator = torch.nn.KLDivLoss()  # KL散度
        elif mi_calculator == "w":
            self.mi_calculator = distance.SinkhornDistance(device=device).to(device)  # Wasserstein距离
        self.temperature =temperature

    def criterion(self,out_dict,y):
        """
        计算总损失
        Args:
            out_dict: 模型输出字典，包含三个分支的预测结果
            y: 真实标签
        Returns:
            return_losses: 包含所有启用的损失项的列表
        """
        return_losses=[]

        # 获取三个分支的输出
        p_y_given_z=out_dict['p_y_given_z']  # 主分类器输出
        p_y_given_f1_f2_f3_f4=out_dict['p_y_given_f_all']  # 全局特征分类器输出
        p_y_given_f1_fn_list=out_dict['p_y_given_f1_fn_list']  # 局部分类器输出列表

        # 1. 计算基础交叉熵损失
        loss_fn = nn.CrossEntropyLoss()
        ce_loss = loss_fn(out_dict['p_y_given_z'], y)

        # 2. 计算全局信息损失
        if self.gil_loss:
            # 全局特征的交叉熵损失
            ce_loss+=loss_fn(p_y_given_f1_f2_f3_f4, y)
            # 计算全局互信息损失（基于KL散度或Wasserstein距离）
            global_mi_loss = self.mi_calculator(
                self.softmax(p_y_given_f1_f2_f3_f4.detach() / self.temperature).log(),
                self.softmax(p_y_given_z / self.temperature)
            )

        # 3. 计算局部信息损失
        if self.lil_loss:
            local_loss = 0
            if self.method == 'distance':
                # 计算不同局部特征之间的距离
                for i in range(len(p_y_given_f1_fn_list)):
                    for j in range(i+1,len(p_y_given_f1_fn_list)):
                        local_loss = local_loss+self.mi_calculator(
                            self.softmax(p_y_given_f1_fn_list[i] / self.temperature).log(),
                            self.softmax(p_y_given_f1_fn_list[j] / self.temperature)
                        )
                local_loss=1-local_loss  # 转换为相似度
            elif self.method == 'mi':
                # 计算局部特征与全局特征之间的互信息
                p_y_given_f1_f2_f3_f4_soft = self.softmax(p_y_given_f1_f2_f3_f4.detach() / self.temperature)
                for out_v in p_y_given_f1_fn_list:
                    local_loss = local_loss + self.mi_calculator(
                        p_y_given_f1_f2_f3_f4_soft.log(),
                        self.softmax(out_v / self.temperature)
                    )
                    ce_loss = ce_loss + loss_fn(out_v, y)  # 每个局部分类器的CE损失
                local_loss = torch.exp(-local_loss)  # 转换为正相关

        # 4. 组合所有损失
        return_losses.append(ce_loss)
        if self.gil_loss:
            return_losses.append(global_mi_loss)
        if self.lil_loss:
            return_losses.append(local_loss)
        return return_losses

    def balance_mult_loss(self,losses):
        """
        平衡多个损失项
        Args:
            losses: 损失列表
        Returns:
            loss: 加权后的总损失
        支持三种平衡方式：
        1. auto: 自动学习权重
        2. hyper: 使用预设权重
        3. 默认: 简单相加
        """
        if self.bml_method == 'auto':
            # 使用自动权重平衡器
            loss =self.balance_loss(losses)
        elif self.bml_method == 'hyper':
            # 使用预设的超参数权重
            loss = 0
            for i, l in enumerate(losses):
                loss = loss+l*self.scales[i]
        else:
            # 直接相加
            loss=sum(losses)
        return loss
